package com.dl.chat.service;

import com.pgvector.PGvector;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

@Service
public class VectorDatabaseService {

    @Autowired
    private JdbcTemplate jdbcTemplate;

    /**
     * 从向量数据库中检索最相关的文档
     *
     * @param queryVector 查询向量
     * @param topK        返回的文档数量
     * @return 相关文档列表
     */
    public List<String> search(float[] queryVector, int topK,float minSimilarity) {
        PGvector pgVector = new PGvector(queryVector);
        //1.计算topK个最相似的文档
        String sql = "SELECT content FROM files ORDER BY embedding <#> ? LIMIT ?";
        List<Map<String,Object>> candidates = jdbcTemplate.queryForList(sql,pgVector);

        //2.计算相似度过滤
        List<String> result = new ArrayList<>();
        for(Map<String,Object> row : candidates){
            float[] dbVector = (float[]) row.get("embedding");
            float similarity = cosineSimilarity(queryVector,dbVector);
            if(similarity > minSimilarity){
                result.add((String) row.get("content"));
            }
        }

        return result;
    }


    // 手动计算余弦相似度（替代 SQL 中的 <=>）
    private float cosineSimilarity(float[] a, float[] b) {
        float dot = 0, normA = 0, normB = 0;
        for (int i = 0; i < a.length; i++) {
            dot += a[i] * b[i];
            normA += a[i] * a[i];
            normB += b[i] * b[i];
        }
        return(float) (dot / (Math.sqrt(normA) * Math.sqrt(normB)));
    }
}


